# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: inkyu

import logging
import os
from typing import List, Optional, Union

import tornado.gen
from tornado.ioloop import PeriodicCallback

import coin.flow.subscriber as flow_subscriber
from coin.base.config import Config
from coin.base.param_util import to_list
from coin.base.timestamp import get_timestamp
from coin.exchange.base.order_container import OrderContainer
from coin.exchange.base.order_gateway import BaseOrderGatewayInfoMutator
from coin.exchange.base.order_gateway_fresh import PositionOrderGatewayFresh, PerProductPositionOrderGatewayFresh
from coin.exchange.base.order_gateway_logger import gen_og_log_request
from coin.exchange.base.order_gateway_ready import SpotOrderGatewayReady
from coin.proto.coin_order_gateway_pb2 import (OrderEvent, OrderGatewayConfig)
from coin.exchange.bitfinex_v2.order_gateway_log import (
    BitfinexV2OrderGatewayLogProvider,)
from coin.exchange.bitfinex_v2.order import BitfinexOrderInternal
from coin.exchange.bitfinex_v2.ws.private_subscriber import (BitfinexV2PrivateSubscriber)
from coin.exchange.base.order_gateway import (
    OrderGatewayBase,
    Order,
    OrderSide,
    OrderType,
    OrderGatewayStatus,
)
from coin.exchange.bitfinex_v2.kr_rest.product import (
    BitfinexCurrency,
    BitfinexProduct,
)
from coin.proto.coin_query_pb2 import (AccountBalance, CurrencyBalance)


def generate_client_order_id(foreign=False):
  timestamp = int(get_timestamp() / 1000)
  generated = int(timestamp % 10**12)
  if foreign:
    return -generated
  else:
    return generated


class BitfinexV2GatewayInfo(BaseOrderGatewayInfoMutator):
  CurrencyType = BitfinexCurrency
  ProductType = BitfinexProduct
  BalanceType = float
  PositionType = None

  def balance_as_proto(self):
    return AccountBalance(exchange='Bitfinex',
                          market_type="Spot",
                          each_balance=self.get_notNone_balances_internal())


class BitfinexV2OrderGateway(OrderGatewayBase):
  _exchange = 'Bitfinex'
  _market_type = 'Spot'
  _api_version = None

  OGInfoType = BitfinexV2GatewayInfo

  def __init__(self,
               currencies: Union[List[BitfinexCurrency], BitfinexCurrency],
               products: Union[List[BitfinexProduct], BitfinexProduct],
               config: Config,
               og_config=None,
               logger=None):
    super().__init__()
    self._logger = logger or logging.getLogger(__name__)
    self._config = config
    self._orders = OrderContainer()

    self._og_ready = SpotOrderGatewayReady()
    self._og_fresh = PerProductPositionOrderGatewayFresh(to_list(products))

    self._og_info = BitfinexV2GatewayInfo(to_list(currencies), to_list(products))
    self._og_logger = BitfinexV2OrderGatewayLogProvider(gen_og_log_request(config, self))

    # Initialized in start.
    self._ioloop = None
    self._flow_sub = None
    self._private_sub = None

    self._og_config = og_config or OrderGatewayConfig()

    self._unacked_order_periodic_cleaner = None
    self._last_private_msg_timestamp = None
    self._balance_periodic_queryer = None

  def get_working_order(self):
    return self._orders.working_orders.values()

  def start(self, queue, ioloop, flow_sub=None):
    self._ioloop = ioloop
    self._private_sub = BitfinexV2PrivateSubscriber(
        self._config, queue, 1, close_callback=self.on_close)  # TODO(inkyu): WORKER_ID
    self._private_sub.start(ioloop)

    self._flow_sub = flow_sub or flow_subscriber.from_queue(queue)
    self._flow_sub.subscribe('bitfinex_v2_ws_private_1', self.on_private_msg)

    # TODO(inkyu): Move me to the base class
    self._unacked_order_periodic_cleaner = PeriodicCallback(
        self.clean_old_unacked_orders, self._og_config.unacked_order_check_period * 1000)
    self._unacked_order_periodic_cleaner.start()

    self._balance_periodic_queryer = PeriodicCallback(self.balance_periodic_query,
                                                      self._og_config.account_query_period * 1000)
    self._balance_periodic_queryer.start()

  def stop(self):
    self._unacked_order_periodic_cleaner.stop()
    self._balance_periodic_queryer.stop()
    self._flow_sub.unsubscribe('bitfinex_v2_ws_private_1', self.on_private_msg)

  @tornado.gen.coroutine
  def _check_private_msg_timestamp(self):
    if self._last_private_msg_timestamp is None:
      return

    time_diff = (get_timestamp() - self._last_private_msg_timestamp) / 1e9
    if time_diff > 30e9:
      self._logger.error('Not receiving private message for %d seconds, reconnect', time_diff)
      self.on_close()
    else:
      self._logger.debug('time_diff = %s', time_diff)

    # for _ in range(30):
    #   if self._og_fresh.check_fresh():
    #     return
    #   else:
    #     yield tornado.gen.sleep(1)

    # if not self._og_fresh.fresh:
    #   self._logger.debug('Not receiving private message for long time, reconnect')
    #   self.on_close()

  def clean_old_unacked_orders(self):
    self._logger.debug('Cleaning old un-acked orders.')
    unacked_order_id_list = self._orders.unacked_order_id_list()
    now_ts = get_timestamp()
    for order_id in unacked_order_id_list:
      order = self._orders.get_working_order(order_id)
      assert not order.accepted, order
      submit_ts = order.internal.timestamp_submit or 0
      if now_ts - submit_ts > 60 * (10**9):  # 60 seconds
        self._og_logger.gen_flush_unacked_order_event_and_log(order.product.symbol,
                                                              order.internal.exchange_order_id,
                                                              now_ts)
        self._orders.remove_from_working_order(order_id)
        self._logger.warning(('Order %s is un-acked more than 60 seconds. '
                              'Removing from on-the-fly order list.') % str(order))
    self._check_private_msg_timestamp()

  def balance_periodic_query(self):
    self._private_sub.request_wallet(
        'exchange', (currency.native_symbol for currency in self._og_info.currencies))

  def _check_order_scope(self, native_symbol, order_type):
    ret = True
    # filter out orders of unsubscribed products and different market
    # type
    product = BitfinexProduct.FromStrNativeProduct(native_symbol)
    if product not in self._og_info.products or \
       order_type != 'EXCHANGE LIMIT':
      ret = False
    return ret

  def _handle_ws_wallet(self, message: list):
    wallet_type = message[0]
    native_currency = message[1]
    balance = float(message[2])
    # unsettled_interest = message[3]
    available = message[4]  # None
    currency = BitfinexCurrency.FromStrNativeCurrencyNothrow(native_currency)
    if currency is None:
      return
    available = available or balance

    if wallet_type == 'exchange':  # exchange, margin, funding
      assert balance >= 0.0, balance
      currency_balance = CurrencyBalance(currency=currency.currency,
                                         currency_native=currency.native_currency,
                                         total=balance,
                                         available=available,
                                         hold=balance - available)
      self._og_info.set_balance(currency, balance, currency_balance)

  def _handle_ws_wallets(self, messages: list):
    self._og_info.set_zero_balances()
    for wallet in messages:
      self._handle_ws_wallet(wallet)
    self._og_logger.write_balance(self._og_info.balance_as_proto())

    if not self.is_ready():
      self._og_ready.set_balance_ready(True)
      self._og_ready.check_ready()
    self._og_fresh.overall().set_balance_timestamp()

  def _handle_ws_order(self, given_order: list):
    exchange_order_id = given_order[0]
    group_id = given_order[1]
    client_order_id = given_order[2]
    if client_order_id is not None:
      client_order_id = int(client_order_id)
    symbol = given_order[3]
    timestamp_created = given_order[4] * 1e6  # milliseconds
    # mts_update = given_order[5]
    qty = given_order[6]  # remaining qty
    qty_orig = given_order[7]
    order_type = given_order[8]  # LIMIT, MARKET, STOP, TRAILING STOP, EXCHANGE MARKET,
    # order_type_prev = given_order[9]  # EXCHANGE LIMIT, EXCHANGE STOP, EXCHANGE TRAILING STOP,
    # flags = given_order[12]
    order_status = given_order[13]  # ACTIVE, EXECUTED, PARTIALLY FILLED, CANCELED
    price_orig = given_order[16]
    price_avg = given_order[17]
    # price_trailing = given_order[18]
    # price_aux_limit = given_order[19]

    product = BitfinexProduct.FromStrNativeProduct(symbol)
    if not self._check_order_scope(symbol, order_type):
      return

    order_side = None
    if qty_orig > 0:
      order_side = OrderSide.BUY
    elif qty_orig < 0:
      order_side = OrderSide.SELL
    else:
      raise ValueError('Zero qty: "%s"' % given_order)

    timestamp = get_timestamp()
    if client_order_id:  # Pass manual order
      order = self._orders.get_working_order(client_order_id)
      if order is None:
        order = self._orders.by_exchange_order_id(exchange_order_id)
        if order and order.order_id < 0:  # Foreign order
          self._orders.remove_from_working_order(order.order_id)
          self._logger.info('Replace foreign order %s -> %s', order.order_id, client_order_id)
          order.order_id = client_order_id
          order.qty = abs(qty_orig)
      if order is None:
        order = Order(order_id=client_order_id,
                      product=product,
                      price=price_orig,
                      qty=abs(qty_orig),
                      order_side=order_side,
                      order_type=OrderType.LIMIT,
                      tag=None,
                      accepted=True,
                      cancel_sent=False,
                      internal=BitfinexOrderInternal.create_by_order_update(
                          price_orig=price_orig,
                          qty_orig=abs(qty_orig),
                          exchange_order_id=exchange_order_id,
                          group_id=group_id,
                          working=True,
                          timestamp=timestamp_created))
      if order.internal.timestamp and order.internal.timestamp != timestamp_created:
        self._logger.warn('%f != %f', order.internal.timestamp, timestamp_created)
      self._orders.add(order)
      order.accepted = True
      order.internal.timestamp = timestamp_created
      order.internal.exchange_order_id = exchange_order_id
      order.internal.group_id = group_id

      self._og_fresh.per(order.product).set_orders_timestamp(timestamp)
      last_fill = order.internal.fill.update_by_remaining_qty(price_avg, abs(qty))
      if (order_status.startswith('CANCELED') or order_status.startswith('POSTONLY CANCELED')
          or order_status.startswith('INSUFFICIENT')):
        # CANCELED was: PARTIALLY FILLED @ 0.0121(0.4), PARTIALLY FILLED @ 0.0121(0.4)
        # INSUFFICIENT BALANCE was: PARTIALLY FILLED @ PRICE(AMOUNT)
        order.internal.working = False
        self._orders.remove_from_working_order(client_order_id)
        self._og_logger.gen_order_event_and_log(order=order,
                                                timestamp=timestamp,
                                                order_event_type=OrderEvent.CANCEL_CONFIRMED)
        if order_status.startswith('INSUFFICIENT'):
          self._logger.error('Insufficient balance: %s', given_order)
      elif order_status.startswith('EXECUTED @'):
        order.internal.working = False
        self._orders.remove_from_working_order(client_order_id)
        self._og_logger.gen_order_event_and_log(order=order,
                                                timestamp=timestamp,
                                                order_event_type=OrderEvent.ORDER_FILLED,
                                                filled_qty=(last_fill.qty if last_fill else 0.))
      elif order_status.startswith('PARTIALLY'):
        self._og_logger.gen_order_event_and_log(order=order,
                                                timestamp=timestamp,
                                                order_event_type=OrderEvent.ORDER_FILLED,
                                                filled_qty=(last_fill.qty if last_fill else 0.))
      elif order_status == 'ACTIVE':  # Active order
        self._og_logger.gen_order_event_and_log(order=order,
                                                timestamp=timestamp,
                                                order_event_type=OrderEvent.ORDER_ACCEPTED)
      elif order_status == 'ERROR':  # Error
        self._logger.error('Unknown order_status: [%s]', given_order)
      else:
        self._logger.error('Unknown order_status: %s', order_status)

      if last_fill is not None:
        self._publish_fill(last_fill.price, last_fill.qty, order=order)
        self._logger.info('%s FILL* (%s, %s): %s %s',
                          order.product.exchange,
                          order.product,
                          order.order_id,
                          order.order_side.name,
                          last_fill.to_str())

  def _handle_ws_order_error(self, message: list, error_message):
    client_order_id = message[2]
    symbol = message[3]
    qty = message[6]
    order_type = message[8]

    if not self._check_order_scope(symbol, order_type):
      return

    timestamp = get_timestamp()
    if client_order_id in self._orders.working_orders:
      self._og_logger.gen_order_event_and_log(order=self._orders.get_working_order(client_order_id),
                                              timestamp=timestamp,
                                              order_event_type=OrderEvent.ORDER_REJECTED)
      self._orders.remove_from_working_order(client_order_id)
    self._logger.error('ORDER ERROR: %s (%s x %s) %s', client_order_id, qty, symbol, error_message)

  def _handle_ws_trade_event(self, message: list):
    # trade_id = message[0]
    pair = message[1]
    # mts_create = message[2]
    exchange_order_id = message[3]
    exec_qty = message[4]
    exec_price = message[5] or 0.
    order_type = message[6]
    order_price = message[7]
    maker = message[8]

    if not self._check_order_scope(pair, order_type):
      return

    if exec_qty > 0:
      order_side = OrderSide.BUY
    elif exec_qty < 0:
      order_side = OrderSide.SELL
    else:
      raise ValueError('Zero qty: "%s"' % str(message))

    timestamp = get_timestamp()
    # TODO(inkyu): Make it log(n)
    order = self._orders.by_exchange_order_id(exchange_order_id)

    if order is None:
      # TODO(inkyu): Remove this
      product = BitfinexProduct.FromStrNativeProduct(pair)
      order_id = generate_client_order_id(foreign=True)
      self._logger.warning('Filled unknown order(%d): %s %f@%f' %
                           (order_id, product, exec_qty, exec_price))

      order = Order(order_id=order_id,
                    product=product,
                    price=order_price,
                    qty=None,
                    order_side=order_side,
                    order_type=OrderType.LIMIT,
                    tag=None,
                    accepted=True,
                    cancel_sent=False,
                    internal=BitfinexOrderInternal.create_by_trade_event(
                        price_orig=None,
                        qty_orig=None,
                        exchange_order_id=exchange_order_id,
                        timestamp=None,
                        maker=maker))
      self._orders.add(order, working=False)
    else:
      order.internal.maker = maker

    # fill message often comes earlier than order status NEW message.
    # update fill
    self._og_fresh.overall().set_fills_timestamp(timestamp)
    self._og_fresh.per(order.product).set_fills_timestamp(timestamp)
    last_fill = order.internal.fill.update_by_last_fill(price_last_fill=exec_price,
                                                        qty_last_fill=abs(exec_qty))
    if last_fill is not None:
      self._publish_fill(last_fill.price, last_fill.qty, order=order)
      self._logger.info('%s FILL+ (%s, %s): %s %s',
                        order.product.exchange,
                        order.product,
                        order.order_id,
                        order.order_side.name,
                        last_fill.to_str())

  def _handle_ws_orders(self, messages: list):
    for given_order in messages:
      self._handle_ws_order(given_order)

    if not self.is_ready():
      self._og_ready.set_orders_ready(True)
      self._og_ready.check_ready()

  def _handle_ws_notification(self, message: list):
    notify_type = message[1]
    notify_info = message[4]  # Payload
    status = message[6]

    if notify_type == 'on-req':  # About the new order
      if status == 'ERROR':
        error_message = message[7]
        self._handle_ws_order_error(notify_info, error_message)
      elif status == 'SUCCESS':
        pass
      elif status == 'FAILURE':  # TODO(inkyu): When this occurs?
        pass
      else:
        self._logger.error('Unknown status: %s (%s)', status, message)
    elif notify_type == 'oc-req':  # About the order cancellation
      if status.startswith('SUCCESS'):
        pass
      elif status.startswith('ERROR'):
        # TODO(inkyu): Separate this.
        exchange_order_id = notify_info[0]
        found_order = self._orders.by_exchange_order_id(exchange_order_id)
        if found_order:
          timestamp = get_timestamp()
          self._og_logger.gen_order_event_and_log(order=found_order,
                                                  timestamp=timestamp,
                                                  order_event_type=OrderEvent.CANCEL_ERROR)
          if message[7] == 'Order not found.':
            if self._orders.remove_from_working_order(found_order.order_id) is not None:
              self._logger.error('Cancel error: %s (%s)', status, found_order)
          else:
            self._logger.error('Cancel error: %s (%s)', status, message)
      else:
        self._logger.error('Unknown cancel status: %s (%s)', status, message)
    elif notify_type == 'oc_multi-req':
      if status == 'INFO':
        pass
      else:
        self._logger.error('Unknown multi-cancel status: %s (%s)', status, message)
    elif notify_type == 'uca':
      self._logger.error('Not implemented notify_type: %s (%s)', notify_type, message)
    elif notify_type == 'fon-req':
      self._logger.error('Not implemented notify_type: %s (%s)', notify_type, message)
    elif notify_type == 'foc-req':
      self._logger.error('Not implemented notify_type: %s (%s)', notify_type, message)
    else:
      self._logger.error('Unknown notify_type: %s (%s)', notify_type, message)

  def on_close(self):
    self._logger.error('Error! Private subscriber is closed! Quit strategy!')
    os._exit(1)

    self._private_sub.stop(trigger_close_callback=False)
    self._og_ready.set_reinit()
    self._og_fresh.reset()
    self._og_info.reset()
    self._og_ready.set_status(OrderGatewayStatus.DISCONNECTED)

    self._logger.error("Bitfinex OG's private subscriber closed, will reconnect soon")
    self._last_private_msg_timestamp = None
    self._private_sub = self._private_sub.clone()  # create a new one.
    self._private_sub.start(self._ioloop)

  def on_private_msg(self, record, queue_data, topic_data):
    self._og_fresh.overall().set_private_timestamp()

    # TODO(inkyu): parse if data is not json
    message = record.data
    self._last_private_msg_timestamp = get_timestamp()

    if type(message) == dict:
      if message['code'] == 20051:  # doc: Stop/Restart Websocket Server (please reconnect)
        self.on_close()
      else:
        raise ValueError('Unknown message: %s' % message)
      return

    event_type = message[1]
    if event_type == 'ps':
      pass
    elif event_type == 'ws':
      self._handle_ws_wallets(message[2])
    elif event_type == 'os':
      self._handle_ws_orders(message[2])
    elif event_type == 'fos':  # Funding offers
      pass
    elif event_type == 'fcs':  # Funding credits
      pass
    elif event_type == 'fls':  # Funding loans
      pass
    elif event_type == 'ats':  # ?
      pass
    elif event_type in ('on', 'ou', 'oc'):
      # New order, Update order, Cancel order. (Involve fill)
      self._handle_ws_order(message[2])
    elif event_type == 'n':  # Notification
      self._handle_ws_notification(message[2])
    elif event_type == 'hb':  # Heartbeat
      # TODO(inkyu): Check heartbeat
      pass
    elif event_type == 'wu':  # Wallet update (Not in the API)
      self._handle_ws_wallet(message[2])
    elif event_type == 'te':  # Trade event
      self._handle_ws_trade_event(message[2])
    elif event_type == 'tu':  # Trade update
      pass
    elif event_type == 'bu':  # Balance info
      pass
    elif event_type == 'miu':  # Margin info
      pass
    elif event_type == 'fiu':  # Funding info
      pass
    else:
      self._logger.error('Unknown event_type: %s (%s)', event_type, message)

  def has_pending_order(self, product):  # TODO(inkyu): Pending => Unacked
    for order in self._orders.working_orders.values():
      if order.product == product and order.internal.exchange_order_id is None:
        return True
    return False

  def submit_impl(
      self,
      product: BitfinexProduct,
      price: float,
      qty: float,
      order_side: OrderSide,
      order_type: OrderType = OrderType.LIMIT,  # LIMIT
      order_id: Optional[int] = None,
      tag=None,
      post_only=False,
      hidden=False):
    assert isinstance(price, float), price
    assert isinstance(qty, float), (qty, qty.__class__)
    assert isinstance(product, BitfinexProduct)
    assert self.is_ready()

    post_qty = None
    if order_side == OrderSide.BUY:
      post_qty = '%f' % qty  # Positive for buy
    elif order_side == OrderSide.SELL:
      post_qty = '-%f' % qty  # Negative for sell
    else:
      raise ValueError('Unknown side: %s' % order_side)
    if order_type != OrderType.LIMIT:
      raise NotImplementedError()

    timestamp = get_timestamp()
    order_id = order_id or generate_client_order_id()
    assert isinstance(order_id, int), ('order_id must be integer type: %s' % str(type(order_id)))
    order = Order(order_id=order_id,
                  product=product,
                  price=price,
                  qty=qty,
                  order_side=order_side,
                  order_type=order_type,
                  tag=None,
                  accepted=False,
                  cancel_sent=False,
                  internal=BitfinexOrderInternal.create_by_submit(timestamp_submit=timestamp,
                                                                  price_orig=price,
                                                                  qty_orig=qty))
    self._orders.add(order)
    self._og_logger.gen_order_request_and_order_event_and_log(order.product.symbol,
                                                              order.price,
                                                              order.qty,
                                                              order.order_side,
                                                              order.order_type,
                                                              str(order.order_id),
                                                              timestamp)
    self._og_logger.gen_order_response_and_order_event_and_log(order.product.symbol,
                                                               True,
                                                               str(order.order_id),
                                                               timestamp)
    return self._private_sub.place_order(product.native_symbol,
                                         None,
                                         str(price),
                                         post_qty,
                                         client_order_id=order.order_id,
                                         post_only=post_only,
                                         hidden=hidden)

  def cancel(self, order_id, *, ignore_error=False):
    # TODO(inkyu): Make order_id to be string
    return self.cancel_impl(order_id, ignore_error=ignore_error)

  def cancel_impl(self, order_id, *, ignore_error=False):
    order_ids = to_list(order_id)
    return self.cancel_multiple(order_ids, ignore_error=ignore_error)

  def cancel_multiple_impl(self, order_ids, *, ignore_error=False):
    assert type(order_ids) == list

    exchange_order_ids = []
    timestamp = get_timestamp()
    for order_id in order_ids:
      order = self._orders.get_working_order(order_id)
      if not order:
        continue
      exchange_order_id = order.internal.exchange_order_id
      if not exchange_order_id:
        self._logger.error('order:%s is not accepted yet.' % order_id)
        continue
      order.cancel_sent = True
      exchange_order_ids.append(exchange_order_id)
      self._og_logger.gen_cancel_request_and_order_event_and_log(order.product.symbol,
                                                                 str(order.order_id),
                                                                 timestamp)
      self._og_logger.gen_cancel_response_and_log(order.product.symbol, True, str(order.order_id))

    if exchange_order_ids:
      return self._private_sub.cancel_by_server_order_ids(exchange_order_ids)

  def cancel_all_impl(self, *, ignore_error=False):
    self.cancel_multiple(list(self._orders.working_orders.keys()))

  def cancel_product_impl(self, product, *, ignore_error):
    order_ids = self.order_ids_by_product(product)
    if order_ids:
      self.cancel_multiple(order_ids)

  def order_ids_by_product(self, product):
    res = []
    for order in self._order.working_orders.values():
      if order.product == product:
        res.append(order.order_id)
    return res
